package cn.lg.soar.common.algorithm;

import javax.crypto.Cipher;
import javax.crypto.SecretKey;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.io.InputStream;
import java.io.OutputStream;
import java.security.MessageDigest;
import java.util.Arrays;

/**
 * 文件编码器
 * @author luguoxiang
 * 开源项目：https://gitee.com/lgx1992/lg-soar 求star！请给我star！请帮我点个star！
 */
public class FileEncoder {

    private static final String ALGORITHM = "AES";
    private static final String TRANSFORMATION = "AES/CBC/PKCS5Padding";

    /**
     * 转Md5
     * @param inputStream 文件流
     * @return
     */
    public static byte[] toMd5(InputStream inputStream) {
        try {
            MessageDigest digest = MessageDigest.getInstance("MD5");
            byte[] buffer = new byte[8192];
            int bytesRead;
            while ((bytesRead = inputStream.read(buffer)) != -1) {
                digest.update(buffer, 0, bytesRead);
            }
            return digest.digest();
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 获取密钥（根据账号密码）
     * @param account 账号
     * @param password 密码
     * @return
     */
    public static Secret getSecret(String account, String password) {
        return new Secret(account, password);
    }

    /**
     * 加密文件
     * @param inputStream 原文件流
     * @param outputStream 加密后的文件流
     * @param secret
     */
    public static void encryptFile(InputStream inputStream, OutputStream outputStream, Secret secret) {
        try {
            doCrypto(Cipher.ENCRYPT_MODE, inputStream, outputStream, secret);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 解密文件
     * @param inputStream 原文件流
     * @param outputStream 加密后的文件流
     * @param secret
     */
    public static void decryptFile(InputStream inputStream, OutputStream outputStream, Secret secret) {
        try {
            doCrypto(Cipher.DECRYPT_MODE, inputStream, outputStream, secret);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 执行加密或解密操作
     * @param cipherMode
     * @param inputStream
     * @param outputStream
     * @param secretKey
     * @throws Exception
     */
    private static void doCrypto(int cipherMode, InputStream inputStream, OutputStream outputStream, Secret secret) throws Exception {
        Cipher cipher = Cipher.getInstance(TRANSFORMATION);
        cipher.init(cipherMode, secret.key, secret.iv);

        byte[] buffer = new byte[16*1024];
        int bytesRead;
        while ((bytesRead = inputStream.read(buffer)) != -1) {
            byte[] output = cipher.update(buffer, 0, bytesRead);
            if (output != null) {
                outputStream.write(output);
            }
        }
        byte[] outputBytes = cipher.doFinal();
        if (outputBytes != null) {
            outputStream.write(outputBytes);
        }
    }

    public static class Secret {

        private final IvParameterSpec iv;

        private final SecretKey key;

        public Secret(IvParameterSpec iv, SecretKey key) {
            this.iv = iv;
            this.key = key;
        }

        public Secret(String account, String password) {
            byte[] bytes = CryptoUtils.sha256(account.getBytes());
            this.iv = new IvParameterSpec(Arrays.copyOfRange(bytes, 0, 16));
            bytes = CryptoUtils.sha256(password.getBytes());
            this.key = new SecretKeySpec(Arrays.copyOfRange(bytes, 0, 16), ALGORITHM);
        }

    }

}